// Tencent is pleased to support the open source community by making ncnn available.
//
// Copyright (C) 2017 THL A29 Limited, a Tencent company. All rights reserved.
//
// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// https://opensource.org/licenses/BSD-3-Clause
//
// Unless required by applicable law or agreed to in writing, software distributed
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.

#include "innerproduct.h"

#include "layer_type.h"

namespace ncnn {

InnerProduct::InnerProduct()
{
    one_blob_only = true;
    support_inplace = false;
}

int InnerProduct::load_param(const ParamDict& pd)
{
    num_output = pd.get(0, 0);
    bias_term = pd.get(1, 0);
    weight_data_size = pd.get(2, 0);
    int8_scale_term = pd.get(8, 0);
    activation_type = pd.get(9, 0);
    activation_params = pd.get(10, Mat());

    if (int8_scale_term)
    {
#if NCNN_INT8
        support_int8_storage = true;
#else
        NCNN_LOGE("please build ncnn with NCNN_INT8 enabled for int8 inference");
        return -1;
#endif
    }

    return 0;
}

/**
 * 全连接层加载权重
 */
int InnerProduct::load_model(const ModelBin& mb)
{
    weight_data = mb.load(weight_data_size, 0);
    if (weight_data.empty())
        return -100;

    if (bias_term)
    {
        bias_data = mb.load(num_output, 1);
        if (bias_data.empty())
            return -100;
    }

#if NCNN_INT8
    // 如果开启了int8量化，则加载维权的缩放系数
    // TODO 缩放系数是如何计算的？
    if (int8_scale_term)
    {
        weight_data_int8_scales = mb.load(num_output, 1);
        bottom_blob_int8_scales = mb.load(1, 1);
    }
#endif // NCNN_INT8

    return 0;
}

/**
 * 全连接层，资源预处理
 * 如果开启了INT8量化处理，那么将会在这个操作里面将
 * 权重量化为INT8格式
 */
int InnerProduct::create_pipeline(const Option& opt)
{
#if NCNN_INT8
    // 开启了INT8量化
    // runtime quantize the weight data
    if (opt.use_int8_inference && weight_data.elemsize == (size_t)4u && int8_scale_term)
    {
        // 输出feature，每个feature权重占用的字节数
        // 根据矩阵的乘法可知，将矩阵转换为w = num_input, h=num_output
        // 这样乘以是可以在矩阵量化中方便操作
        const int num_input = weight_data_size / num_output;

        // 将模型权重进行重排序，按照
        // 方便后续全连接乘法操作
        Mat weight_data_r2 = weight_data.reshape(num_input, num_output);

        Mat weight_data_int8;
        Option opt_q = opt;
        opt_q.use_packing_layout = false;
        // 将矩阵进行INT8量化，得到量化后的矩阵
        quantize_to_int8(weight_data_r2, weight_data_int8, weight_data_int8_scales, opt_q);
        if (weight_data_int8.empty())
            return -100;

        // 矩阵重排序，将权重矩阵展平
        weight_data = weight_data_int8.reshape(weight_data_size);
    }
#endif // NCNN_INT8

    return 0;
}

/**
 * 全连接层的推理，整理方法就是通过矩阵运算
 * 优化点：
 * 1. 多线程多样本并行计算
 * 2. 权重矩阵中每一列按行存储，保证计算时可以一次性取出所需的列进行计算
 * 
 * @param bottom_blob:输入矩阵
 * @param top_blob:输出矩阵
 * @param opt：层推理参数
 * 
 * @return: int，返回-100表示创建输出矩阵失败, 返回0则表示全连接层推理成功
 */
int InnerProduct::forward(const Mat& bottom_blob, Mat& top_blob, const Option& opt) const
{
#if NCNN_INT8
    // 如果开启了INT8量化，则执行量化推理
    if (opt.use_int8_inference && weight_data.elemsize == (size_t)1u)
    {
        // todo 看到这里
        return forward_int8(bottom_blob, top_blob, opt);
    }
#endif

    // 计算全连接层的输入尺寸
    const int num_input = weight_data_size / num_output;

    int w = bottom_blob.w;
    int h = bottom_blob.h;
    int channels = bottom_blob.c;
    size_t elemsize = bottom_blob.elemsize;
    int size = w * h;
    
    if (bottom_blob.dims == 2 && w == num_input && h > 1)
    {
        // 输入矩阵的维度=2，及仅有w h两维
        // 输入矩阵的宽度正好等于全连接层的输入尺寸
        // 且输入矩阵的高度（第二维）尺寸不是1，即输入矩阵不是shape=(None, 1)这种形式的矩阵，这种形式的矩阵表示输入样本的数量不为1

        // gemm
        // 创建一个输出矩阵，其第一维等于全连接层的输出尺寸，表示经过转换后的尺寸
        // 而输出矩阵构建的第二个参数h表示有多少个训练样本
        top_blob.create(num_output, h, elemsize, opt.blob_allocator);
        if (top_blob.empty())
            return -100;

        // 使用openmp开启多线程，每个训练样本并行训练
        #pragma omp parallel for num_threads(opt.num_threads)
        for (int j = 0; j < h; j++)
        {
            // m 输入矩阵的第j个样本
            const float* m = bottom_blob.row(j);
            // outptr 存储经过全连接层后，输出样本的存储缓存
            float* outptr = top_blob.row(j);
            
            // 以输出维度进行遍历，计算每一个输出节点的值
            // 根据矩阵的计算方法，weight_data的尺寸形如(w: bootom_blob width 即输入矩阵的宽度 num_input, h: num_output)
            // 这里weight_data的存储顺序是每一行存储的是在计算矩阵时的每一列
            for (int p = 0; p < num_output; p++)
            {
                // 提取需要计算的那一列（在缓存中，由于这一列是按行连续存储，所以 weight_data（首地址） + w * p（第N个输出维度节点））
                const float* kptr = (const float*)weight_data + w * p;

                // 存储输出节点的值
                float sum = 0.f;

                // 如果存在偏置，则将偏置提前赋值到sum
                if (bias_term)
                    sum = bias_data[p];

                // 按照矩阵的乘法，输入样本的每一个值 * 权重矩阵的当前列的每一个值
                // 将计算出来的结果保存到sum中
                for (int i = 0; i < w; i++)
                {
                    sum += m[i] * kptr[i];
                }

                // 根据设置的激励函数，进行相应的激励函数的计算
                if (activation_type == 1)
                {   
                    sum = std::max(sum, 0.f);
                }
                else if (activation_type == 2)
                {
                    float slope = activation_params[0];
                    sum = sum > 0.f ? sum : sum * slope;
                }
                else if (activation_type == 3)
                {
                    float min = activation_params[0];
                    float max = activation_params[1];
                    if (sum < min)
                        sum = min;
                    if (sum > max)
                        sum = max;
                }
                else if (activation_type == 4)
                {
                    sum = static_cast<float>(1.f / (1.f + exp(-sum)));
                }
                else if (activation_type == 5)
                {
                    sum = static_cast<float>(sum * tanh(log(exp(sum) + 1.f)));
                }

                // 将计算的结果保存到输出节点
                outptr[p] = sum;
            }
        }

        return 0;
    }

    // 如果输入矩阵的维度大于2则进入下面的推理流程
    // 创建输出矩阵
    // 输出矩阵的宽度是全连接层的输出维度
    top_blob.create(num_output, elemsize, opt.blob_allocator);
    if (top_blob.empty())
        return -100;

    // num_output
    // 这里多线程就按照配置的线程数进行
    // 这里需要搞清楚，多维度的输入样本的表示、计算方式
    // 而在高维矩阵计算方式中，输出矩阵topBlob的 shape is （1，num_output）
    // 那么输入矩阵的shape is (1, channels, w, h)
    #pragma omp parallel for num_threads(opt.num_threads)
    for (int p = 0; p < num_output; p++)
    {
        //sum 输出维度每个点的输出值
        float sum = 0.f;

        // 存在偏置则赋值偏置
        if (bias_term)
            sum = bias_data[p];

        // channels
        // 如果每个样本的维度是高维的，还具备通道数，则需要以通道数遍历计算输出节点的值
        // 对于高维样本，则权重也是高维，多通道的
        // 根据以下的算法，猜测权重维度的排列是（num_output, channels, w, h）
        // 那么高维矩阵的全连接层，其计算方式是每个通道都与输入的样本每个通道进行计算
        // bottom_blob的shape is (channels, w, h)
        // todo 高维矩阵的全连接层，输入的数据含义，计算含义
        for (int q = 0; q < channels; q++)
        {
            // size * channels * p: 由于权重维度的排列方式，那么这个计算式代表是：找到当前输出点的权重参数
            // size * q: 找到当前输出权重，当前通道的首地址
            const float* w = (const float*)weight_data + size * channels * p + size * q;
            // 获取输入矩阵当前通道的样本首地址
            const float* m = bottom_blob.channel(q);

            // 权重与样本，矩阵对应位置乘积
            for (int i = 0; i < size; i++)
            {
                sum += m[i] * w[i];
            }
        }

        // 根据设置的激活函数进行相应的计算
        if (activation_type == 1)
        {
            sum = std::max(sum, 0.f);
        }
        else if (activation_type == 2)
        {
            float slope = activation_params[0];
            sum = sum > 0.f ? sum : sum * slope;
        }
        else if (activation_type == 3)
        {
            float min = activation_params[0];
            float max = activation_params[1];
            if (sum < min)
                sum = min;
            if (sum > max)
                sum = max;
        }
        else if (activation_type == 4)
        {
            sum = static_cast<float>(1.f / (1.f + exp(-sum)));
        }
        else if (activation_type == 5)
        {
            sum = static_cast<float>(sum * tanh(log(exp(sum) + 1.f)));
        }

        // 将输出节点得到的值存放在输出矩阵上
        top_blob[p] = sum;
    }

    return 0;
}

#if NCNN_INT8
/**
 * int8 量化 推理
 * 
 * @return： 返回-100则表示创建失败，返回
 */
int InnerProduct::forward_int8(const Mat& bottom_blob, Mat& top_blob, const Option& opt) const
{
    // 计算输入样本的维度
    const int num_input = weight_data_size / num_output;

    // 输入矩阵，输入样本的维度
    int w = bottom_blob.w;
    // 输入矩阵，输入样本的个数
    int h = bottom_blob.h;
    // 输入矩阵的通道数
    int channels = bottom_blob.c;
    // 输入矩阵的单个像素的字节数
    size_t elemsize = bottom_blob.elemsize;
    // 计算通道的尺寸
    int size = w * h;

    // 将矩阵的进行int8量化
    Mat bottom_blob_int8 = bottom_blob;
    // 如果elemsize的单个元素的字节数不是一个字节，则需要进行int8量化
    // 为啥？elemsize内部定义的
    if (elemsize != 1)
    {
        // todo 这里use_packing_layout优化是关闭的？为啥？
        Option opt_g = opt;
        opt_g.blob_allocator = opt.workspace_allocator;
        opt_g.use_packing_layout = false;

        quantize_to_int8(bottom_blob, bottom_blob_int8, bottom_blob_int8_scales, opt_g);
    }

    // 这里和非量化矩阵的全连接层的推理一样，根据维度进行区分计算
    // todo 看到这里 
    // 虽然整体一致，但是由于细节上存在多处不一致，强行融合在一起会导致
    // 代码整体复杂，所以分开处理虽然代码重复，但是结构流程清晰
    if (bottom_blob.dims == 2 && w == num_input && h > 1)
    {
        // gemm
        // 创建输出矩阵，唯一不同的是，量化矩阵其输出矩阵
        // 的单个元素像素为4个字节，也就是float
        // num_output代表输出节点的个数，也就是每个样本经过全连接层后输出的数据
        // h代表有多少个样本，就存储多少个输出数据
        top_blob.create(num_output, h, 4u, opt.blob_allocator);
        if (top_blob.empty())
            return -100;

        // 一致
        #pragma omp parallel for num_threads(opt.num_threads)
        for (int j = 0; j < h; j++)
        {
            // 获取输入量化矩阵当前的输入样本
            // 与非量化版本不同之处在于，其首地址是char类型
            const signed char* m = bottom_blob_int8.row<signed char>(j);
            // 获取输出矩阵当前样本的输出缓存首地址
            float* outptr = top_blob.row(j);

            // 遍历每一个全连接层的输出节点
            for (int p = 0; p < num_output; p++)
            {
                // 获取权重首地址，从这里可以看出，weight_data的shape is （num_output， w）
                // 而创建矩阵时采用的shape is (w（width）, num_output（高）)
                // kptr获取当前输出节点对应的权重首地址
                const signed char* kptr = (const signed char*)weight_data + w * p;
                int sum = 0;

                // todo 为什么量化层的偏置要放在后面

                // 权重与输入矩阵对应样本数据相乘
                for (int i = 0; i < w; i++)
                {
                    sum += m[i] * kptr[i];
                }
                
                // 经过矩阵乘法后，将量化后的矩阵反转换为float类型
                // dequantize and relu
                float scale_in;
                // todo 以下两个条件的计算方式是什么意义
                if (weight_data_int8_scales[p] == 0)
                    scale_in = 0;
                else
                    scale_in = 1.f / (bottom_blob_int8_scales[0] * weight_data_int8_scales[p]);

                // 将矩阵乘法后计算得到的sum 反量化为 float32类型
                float sumfp32 = sum * scale_in;

                // 计算偏置，因为偏置是针对非量化的矩阵乘积结果
                // 所以需要在反量化后才能进行偏置计算
                if (bias_term)
                    sumfp32 += bias_data[p];

                // 根据激活函数进行计算
                if (activation_type == 1)
                {
                    sumfp32 = std::max(sumfp32, 0.f);
                }
                else if (activation_type == 2)
                {
                    float slope = activation_params[0];
                    sumfp32 = sumfp32 > 0.f ? sumfp32 : sumfp32 * slope;
                }
                else if (activation_type == 3)
                {
                    float min = activation_params[0];
                    float max = activation_params[1];
                    if (sumfp32 < min)
                        sumfp32 = min;
                    if (sumfp32 > max)
                        sumfp32 = max;
                }
                else if (activation_type == 4)
                {
                    sumfp32 = static_cast<float>(1.f / (1.f + exp(-sumfp32)));
                }
                else if (activation_type == 5)
                {
                    sumfp32 = static_cast<float>(sumfp32 * tanh(log(exp(sumfp32) + 1.f)));
                }

                // 将计算的结果存储在输出缓存上
                outptr[p] = sumfp32;
            }
        }

        return 0;
    }

    // 针对高维输入的样本，创建浮点类型的输出矩阵
    // 而且这个矩阵的shape is （1. num_output）
    // todo 从以下代码，感觉高维矩阵只有一个输入样本，那么高维矩阵的意义是什么？
    // 用于什么场景下？
    top_blob.create(num_output, 4u, opt.blob_allocator);
    if (top_blob.empty())
        return -100;

    // num_output
    #pragma omp parallel for num_threads(opt.num_threads)
    for (int p = 0; p < num_output; p++)
    {
        // 获取输出矩阵的首地址
        float* outptr = top_blob;

        int sum = 0;

        // 高维输入样本矩阵下，其weihgt_data的shape is（num_output, channels, w, h）
        // 所以这个offset是找到当前遍历的输出节点的权重矩阵
        int offset = size * channels * p;
        // channels
        for (int q = 0; q < channels; q++)
        {
            // 根据通道数 获取输出节点的第n个通道的首地址
            const signed char* w = (const signed char*)weight_data + offset + size * q;
            // 量化后的高维输入矩阵，当前通道首地址
            // bottom_blob_int8 shape is (1, channels, w, h)
            const signed char* m = bottom_blob_int8.channel(q);

            for (int i = 0; i < size; i++)
            {
                sum += m[i] * w[i];
            }
        }

        // dequantize and relu
        // 和之前一样，矩阵计算完毕后进行反量化
        // 以下过程均一致
        float scale_in;
        if (weight_data_int8_scales[p] == 0)
            scale_in = 0;
        else
            scale_in = 1.f / (bottom_blob_int8_scales[0] * weight_data_int8_scales[p]);

        float sumfp32 = sum * scale_in;

        if (bias_term)
            sumfp32 += bias_data[p];

        if (activation_type == 1)
        {
            sumfp32 = std::max(sumfp32, 0.f);
        }
        else if (activation_type == 2)
        {
            float slope = activation_params[0];
            sumfp32 = sumfp32 > 0.f ? sumfp32 : sumfp32 * slope;
        }
        else if (activation_type == 3)
        {
            float min = activation_params[0];
            float max = activation_params[1];
            if (sumfp32 < min)
                sumfp32 = min;
            if (sumfp32 > max)
                sumfp32 = max;
        }
        else if (activation_type == 4)
        {
            sumfp32 = static_cast<float>(1.f / (1.f + exp(-sumfp32)));
        }
        else if (activation_type == 5)
        {
            sumfp32 = static_cast<float>(sumfp32 * tanh(log(exp(sumfp32) + 1.f)));
        }

        outptr[p] = sumfp32;
    }

    return 0;
}
#endif // NCNN_INT8

} // namespace ncnn
